package tests;

import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedList;

import utils.StateCamera;
import utils.StateCamera.CameraType;
import utils.math_utils.Coordinate;
import utils.math_utils.MathUtils;
import utils.math_utils.RandStream;
import Jama.Matrix;

import components.AbstractNeuron;
import components.LIFNeuron;
import components.Network;
import components.Synapse;

import connectors.GridConnect;

public class SeparationTest {

	public static final String testFileName = "ROCCurves.dat";
	public static final int firstPass = 2100;
	public static final int trainEpochs = 2100;
	public static final int testEpochs = 999;
	public static final double timeStep = 0.1; //ms
	public static double time = 0;
	public static final double pulseDuration = 0.3;//ms
	public static final int epoch = 30; //ms	
	public static final int numRes = 1024;
	public static final int memTc = 30;

	private StateCamera camera;
	private static int [] resDims = {16, 16, 4};
	private final Network net = new Network();
	private static ArrayList<boolean []> trainingIn;
	private static ArrayList<boolean []> testingIn;	
	private LIFNeuron [] res = new LIFNeuron[numRes];
	private static int stimGrpSize = resDims[2];
	private LIFNeuron [] stimGrp1 = new LIFNeuron[stimGrpSize];
	private LIFNeuron [] stimGrp2 = new LIFNeuron[stimGrpSize];
	private LIFNeuron [] stimGrp3 = new LIFNeuron[stimGrpSize];
	private static double [][] runningAverageState1 =
			new double[epoch][numRes];
	private static double [][] runningAverageState2 =
			new double[epoch][numRes];
	private static double [][] runningAverageState3 =
			new double[epoch][numRes];
	private static byte [][][] stateMatrix;
	private static Matrix [] filters = new Matrix[epoch];

	
	public SeparationTest(){
		initialize();
		generateInputSets();
	}
	
	private void generateInputSets(){
		boolean [][] binaryMap = new boolean [3][3];
		binaryMap[0] = new boolean[]{false, true, true};
		binaryMap[1] =  new boolean[]{true, false, true};
		binaryMap[2] = new boolean[]{true, true, false};	
		trainingIn = new ArrayList<boolean []>(trainEpochs);
		testingIn = new ArrayList<boolean []>(999);	
		for(int i = 0; i < testEpochs/3; i++) {
			testingIn.add(binaryMap[0]);
			testingIn.add(binaryMap[1]);
			testingIn.add(binaryMap[2]);		
		}
		for(int i = 0; i < firstPass/3; i++) {
			trainingIn.add(binaryMap[0]);
			trainingIn.add(binaryMap[1]);
			trainingIn.add(binaryMap[2]);	
		}		
		Collections.shuffle(testingIn);
		Collections.shuffle(trainingIn);	
		
		
	}
	
	/**
	 * 
	 */
	private void initialize() {
			
		camera = new StateCamera(CameraType.Bin,
				timeStep, 1);			
		
		res =(Network.initializeNeurons(numRes,
				memTc, 15, 2.5, 0.2, camera));
	
		net.cubicLayout(new Coordinate(0, 0, 0), resDims,
				res);
		
		int x1 = 4;
		int x2 = 8;
		int x3 = 12;
		int y = 6;	
		
		HashMap<Coordinate, AbstractNeuron> map = net.getNeuronMap();
		int ind1 = 0;
		int ind2 = 0;
		int ind3 = 0;
		
		for(int z = 0; z < resDims[2]; z++) {
			
			stimGrp1[ind1++] = (LIFNeuron) map.get(new Coordinate(x1, y, z));
//			stimGrp1[ind1++] = (LIFNeuron) map.get(new Coordinate(x1 + 1, y, z));
//			stimGrp1[ind1++] = (LIFNeuron) map.get(new Coordinate(x1 - 1, y, z));
//			stimGrp1[ind1++] = (LIFNeuron) map.get(new Coordinate(x1, y - 1, z));
//			stimGrp1[ind1++] = (LIFNeuron) map.get(new Coordinate(x1, y + 1, z));
			
			stimGrp2[ind2++] = (LIFNeuron) map.get(new Coordinate(x2, y, z));
//			stimGrp2[ind2++] = (LIFNeuron) map.get(new Coordinate(x2 + 1, y, z));
//			stimGrp2[ind2++] = (LIFNeuron) map.get(new Coordinate(x2 - 1, y, z));
//			stimGrp2[ind2++] = (LIFNeuron) map.get(new Coordinate(x2, y - 1, z));
//			stimGrp2[ind2++] = (LIFNeuron) map.get(new Coordinate(x2, y + 1, z));
			
			stimGrp3[ind3++] = (LIFNeuron) map.get(new Coordinate(x3, y, z));
//			stimGrp3[ind3++] = (LIFNeuron) map.get(new Coordinate(x3 + 1, y, z));
//			stimGrp3[ind3++] = (LIFNeuron) map.get(new Coordinate(x3 - 1, y, z));
//			stimGrp3[ind3++] = (LIFNeuron) map.get(new Coordinate(x3, y - 1, z));
//			stimGrp3[ind3++] = (LIFNeuron) map.get(new Coordinate(x3, y + 1, z));
			
		}	
		
		GridConnect gCon = new GridConnect();
		Network.initializeSynRandomizers(gCon);
		Synapse [][] syns = gCon.stochasticConnect(res, res, 3.0);
		int counter = 0;
		for(int i = 0; i < syns.length; i++){
			for(int j = 0; j < syns[0].length; j++){
				if(syns[i][j] != null){
					counter++;
				}
			}
		}
		System.out.println("Sparsity: " + (double)counter/(1024*1024));
		
		for(LIFNeuron n : res) {
			n.setStdpOn(false);
			n.setReporter(false);
			n.setRecorder(camera);
			//Homeostatic plasticity ON
			//n.setHmP(false);
		}
		
	}
	
	/**
	 * 
	 * @param test
	 * @param inputs
	 * @param recording
	 */
	public static void run (SeparationTest test, ArrayList<boolean[]> inputs,
			boolean recording, boolean testing) {
		
		int num;
		String process;
		if(!recording && !testing){
			num = firstPass;
			process = "First Pass";
		} else if (recording && !testing) {
			num = trainEpochs;
			process = "Training";
		} else {
			num = testEpochs;
			process = "testing";
		}
		System.out.println(process);
		
		stateMatrix  = new byte[num][epoch][numRes];		
		StateCamera camera = test.getCamera();
		LinkedList<LIFNeuron> spikingNeurons;
		boolean a;
		boolean b;
		boolean c;
		boolean [] current;
		double per = 0.1;
		double runningPer = 0;
		for(int i = 0, m = inputs.size(); i < m; i++) {
			current = inputs.get(i);
			a = current[0] && current[1];
			b = current[0] && current[2];
			c = current[1] && current[2];
			
			if(current[0]){
				for(LIFNeuron n : test.getStimGrp1()) {
					n.setBgCurrent(n.getBgCurrent() + 5000);
				}
			}
			if(current[1]){
				for(LIFNeuron n : test.getStimGrp2()) {
					n.setBgCurrent(n.getBgCurrent() + 5000);
				}
			}
			if(current[2]){
				for(LIFNeuron n : test.getStimGrp3()) {
					n.setBgCurrent(n.getBgCurrent() + 5000);
				}
			}		
			
			if(recording){
				camera.readyFilm(epoch, numRes);
//				//RESETTING VALUES
//				test.reset();
			}
			double t = 0;
			while(t < epoch) {
				
				if(t > pulseDuration){
					for(int j = 0; j < stimGrpSize; j++){
						test.getStimGrp1()[j].setBgCurrent(0.0);
						test.getStimGrp2()[j].setBgCurrent(0.0);
						test.getStimGrp3()[j].setBgCurrent(0.0);
					}
				}
				
				spikingNeurons = new LinkedList<LIFNeuron>();
				for(LIFNeuron n : test.getRes()) {
					if(n.partialUpdate(timeStep)){
						spikingNeurons.add(n);
					}
				}
				
				while(!spikingNeurons.isEmpty()){
					spikingNeurons.poll().spike(timeStep);
				}
				
				if(recording)
					camera.collectFrame();
				
				t += timeStep;
				time += timeStep;
				
			}
			
			runningPer = (double)i/m;
			if(runningPer >= per) {
				per += 0.1;
				System.out.print(Math.floor(runningPer * 100) + "% ");
			}
			
			if(recording){
				stateMatrix[i] = MathUtils.toByte(camera.getCapturedStates());
				if(a){
					runningAverageState1 = MathUtils.arrayAdd2D(
							runningAverageState1, camera.getCapturedStates());
				} else if(b) {
					runningAverageState2 = MathUtils.arrayAdd2D(
							runningAverageState2, camera.getCapturedStates());
				} else if (c) {
					runningAverageState3 = MathUtils.arrayAdd2D(
							runningAverageState3, camera.getCapturedStates());
				} else {
					throw new IllegalStateException("Invalid inputs");
				}
			}	
		}
		
		if(recording){
			MathUtils.arrayScaleInPlace2D(runningAverageState1,
					1.0/((double)inputs.size()/3.0));
			MathUtils.arrayScaleInPlace2D(runningAverageState2,
					1.0/((double)inputs.size()/3.0));
			MathUtils.arrayScaleInPlace2D(runningAverageState3,
					1.0/((double)inputs.size()/3.0));
			
			printToFile("RunningAvg1.dat", runningAverageState1);
			printToFile("RunningAvg2.dat", runningAverageState2);
			printToFile("RunningAvg3.dat", runningAverageState3);
			
			if(testing) {
				test.test(stateMatrix, inputs);
			} else {
				test.train(stateMatrix, inputs);
			}
		}
		
		
		System.out.println();
	}
	
	/**
	 * TRAINNG
	 * @param states
	 * @param inputs
	 */
	public void train(byte [][][] states, ArrayList<boolean[]> inputs) {
		System.out.println();
		Matrix S;
		Matrix dOuts = new Matrix(convert(inputs));
		double [][] sm;
		for(int i = 0; i < epoch; i++) {
			sm = new double[states.length][numRes];
			for(int j = 0; j < states.length; j++) {
				sm[j] = MathUtils.toDouble(states[j][i]);
			}
			MathUtils.addNoise(new RandStream(-0.0005, 0.0005), sm);
			S = new Matrix(sm);
			try {
			filters[i] = S.solve(dOuts);
			} catch (RuntimeException re) {
				System.out.println(re.getMessage() + " Time: " + i);
				filters[i] = null;
				continue;
			}
		}		
	}
	
	/**
	 * TEST
	 * @param states
	 * @param inputs
	 */
	public void test(byte [][][] states, ArrayList<boolean []> inputs) {
		Matrix S;
		double [][] dOuts = convert(inputs);
		double [][] sm;
		try {
			FileWriter fw = new FileWriter(testFileName);
			PrintWriter pw = new PrintWriter(fw);
			System.out.println();
			System.out.print("[");
		for(int i = 0; i < epoch; i++) {
			if(filters[i] != null){
				sm = new double[states.length][numRes];
				for(int j = 0; j < states.length; j++) {
					sm[j] = MathUtils.toDouble(states[j][i]);
				}
				S = new Matrix(sm);
				double [][] outs = (S.times(filters[i])).getArray();
				
				double [][] roc = ROCCurve(-0.5, 0.1, outs, dOuts);
				pw.println(Arrays.toString(roc[0]));
				pw.println(Arrays.toString(roc[1]));
				pw.println();
			} else {
				pw.println(Arrays.toString(new double [100]));
				pw.println(Arrays.toString(new double [100]));
			}
		}
		System.out.print("]");
		System.out.println();
		pw.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	
	/**
	 * Creates an ROC curve from the data given a starting seed (used for
	 * finding the highest threshold for which true positives and false
	 * alarms are 100%) and an interval.
	 * @param seed
	 * @param interval
	 * @param outs
	 * @param targetOuts
	 * @return
	 */
	public double [][] ROCCurve(double seed, double interval,
			double [][] outs, double [][] targetOuts) {
		double [] tpfa;
		double tStart = findROCStart(seed, interval, outs, targetOuts);
		ArrayList<double []> roc = new ArrayList<double []>();
		double threshold = tStart;
		do {
			tpfa = ROC(threshold, outs, targetOuts);
			roc.add(Arrays.copyOf(tpfa, 2));
			threshold += interval;
		}while(tpfa[0] > 0 || tpfa[1] > 0);
		
		double [][] curve = new double[2][roc.size()];
		int i = 0;
		double auc = 0;
		for(double [] d : roc){
			curve[0][i] = d[0];
			curve[1][i] = d[1];
			if(i > 0){
				auc += (curve[1][i]-curve[1][i-1]) * curve[0][i-1];
			}
			i++;
		}
		System.out.print(auc + " ");
		return curve;
	}
	
	/**
	 * Finds the highest threshold for which true positives and false alarms
	 * are both 100%. 
	 * @param seed
	 * @param interval
	 * @param outs
	 * @param targetOuts
	 * @return
	 */
	private double findROCStart(double seed, double interval, double[][] outs,
			double [][] targetOuts) {
		double [] tpfa = ROC(seed, outs, targetOuts);
		double threshold = seed;
		if(tpfa[0] == 1 && tpfa[1] == 1) {
			do {
				threshold += interval;
				tpfa = ROC(threshold, outs, targetOuts);		
			}while(tpfa[0] == 1 && tpfa[1] == 1); 
			return threshold - interval;
		} else {
			do {
				threshold -= interval;
				tpfa = ROC(threshold, outs, targetOuts);		
			}while(tpfa[0] < 1 || tpfa[1] < 1); 
			return threshold;
		}		
	}
	
	/**
	 * Finds the ROC (Receiver Operator Characteristic: false alarm % vs.
	 * true positive %) for a given threshold given some outputs and desired
	 * outputs. 
	 * @param threshold
	 * @param outs
	 * @param targetOuts
	 * @return
	 */
	private double [] ROC(double threshold, double [][] outs,
			double [][] targetOuts) {
		
		double tp = 0;
		double fa = 0;	
		for(int i = 0; i < outs.length; i++){
			for(int j = 0; j < outs[0].length; j++) {
				int val = outs[i][j] >= threshold ? 1 : 0;
				if(val == 1) {
					if(targetOuts[i][j] == 1) {
						tp++;
					} else {
						fa++;
					}
				}	
			}
		}
		
		return new double [] {tp/(outs.length * 2), fa/outs.length};
		
		
		
	}
	
	public double [][] convert(ArrayList<boolean []> inputs) {
		double [][] ins = new double[inputs.size()][3];
		
		for(int i = 0; i < inputs.size(); i++){
			for(int j = 0; j < 3; j++){
				ins[i][j] = inputs.get(i)[j] ? 1.0 : 0.0;
			}
		}		
		return ins;	
	}
	
	/**
	 * RESET
	 */
	public void reset(){
		for(LIFNeuron n : res) {
			n.setMemPotential(RandStream.nextUniform(13.5, 15.1));
			n.setLastSpike(0);
			n.setTime(0);
			for(Synapse s : n.getFanIn()) {
				s.setTime(0);
				s.clearSpkQueues();
			}
		}
	}
	
	/**
	 * 
	 * @param filename
	 * @param mat
	 */
	public static void printToFile(String filename, double [][] mat) {
		try {
			FileWriter fw = new FileWriter(filename);
			PrintWriter pw = new PrintWriter(fw);
			Matrix m = new Matrix(mat);
			
			m.print(pw, 3, 4);
			
			pw.close();
			
		} catch (IOException ie) {
			ie.printStackTrace();
		}
	}
	
	
	/**
	 * @param args
	 */
	public static void main(String[] args) {
		
		SeparationTest test = new SeparationTest();
		
		/*
		 * FIRST PASS
		 */
		run(test, trainingIn, false, false);
		
		for(LIFNeuron n : test.getRes()) {
			n.setReporter(true);
		}	
		
		/*
		 * TRAINING & SEPARATION RECORDING
		 */
		run(test, trainingIn, true, false);
		
		double [] runningDist = new double [epoch];		
		for(int i = 0; i < epoch; i++) {
			runningDist[i] += MathUtils.euclideanDist(runningAverageState1[i],
					runningAverageState2[i]);
			runningDist[i] += MathUtils.euclideanDist(runningAverageState1[i],
					runningAverageState3[i]);
			runningDist[i] += MathUtils.euclideanDist(runningAverageState2[i],
					runningAverageState3[i]);
			runningDist[i] /= 3;
		}	
		System.out.println(Arrays.toString(runningDist));
		
		/*
		 * TESTING
		 */
		run(test, testingIn, true, true);	
	}
	
	
	/*/*******************************************************************
	 * 						GETTERS AND SETTERS							 *
	 *********************************************************************/
	
	public StateCamera getCamera() {
		return camera;
	}

	public void setCamera(StateCamera camera) {
		this.camera = camera;
	}

	public static int[] getResDims() {
		return resDims;
	}

	public static void setResDims(int[] resDims) {
		SeparationTest.resDims = resDims;
	}

	public ArrayList<boolean[]> getTrainingIn() {
		return trainingIn;
	}

	public ArrayList<boolean[]> getTestingIn() {
		return testingIn;
	}

	public LIFNeuron[] getRes() {
		return res;
	}

	public void setRes(LIFNeuron[] res) {
		this.res = res;
	}

	public LIFNeuron[] getStimGrp1() {
		return stimGrp1;
	}

	public void setStimGrp1(LIFNeuron[] stimGrp1) {
		this.stimGrp1 = stimGrp1;
	}

	public LIFNeuron[] getStimGrp2() {
		return stimGrp2;
	}

	public void setStimGrp2(LIFNeuron[] stimGrp2) {
		this.stimGrp2 = stimGrp2;
	}

	public LIFNeuron[] getStimGrp3() {
		return stimGrp3;
	}

	public void setStimGrp3(LIFNeuron[] stimGrp3) {
		this.stimGrp3 = stimGrp3;
	}

	public static int getFirstpass() {
		return firstPass;
	}

	public static int getTestepochs() {
		return testEpochs;
	}

	public static double getTimestep() {
		return timeStep;
	}

	public static double getEpoch() {
		return epoch;
	}

	public static int getNumres() {
		return numRes;
	}

	public static int getMemtc() {
		return memTc;
	}

	public Network getNet() {
		return net;
	}

}
